import os
import cv2
import shutil
from rich.progress import Progress

os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"

src = os.path.dirname(os.path.abspath(__file__))
if not src.endswith(os.path.sep):
    src += os.path.sep


files = os.listdir(src)
total = len(files)

def get_file_path():
    directory = fr"{src}"
    filenames = os.listdir(directory)
    image_roka = [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".gif"]
    return [os.path.join(directory, filename) for filename in filenames
            if os.path.splitext(filename)[1].lower() in image_roka]

def get_file():
    directory = fr"{src}"
    filenames = os.listdir(directory)
    image_roka = [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".gif"]
    return [filename for filename in filenames
            if os.path.splitext(filename)[1].lower() in image_roka]
images = get_file_path()
image_names = get_file()


os.makedirs(os.path.join(src, "Transverse"), exist_ok=True)
os.makedirs(os.path.join(src, "Longitudinal"), exist_ok=True)
os.makedirs(os.path.join(src, "Square"), exist_ok=True)


with Progress() as progress:
    task = progress.add_task("[magenta]处理中...[/magenta]", total=total)
    for jsq in range(len(images)):
        img_data = cv2.imread(images[jsq])
        if img_data is not None:
            height, width, channels = img_data.shape
            if width > height:
                dst_dir = os.path.join(src, "Transverse")
            elif height > width:
                dst_dir = os.path.join(src, "Longitudinal")
            else:
                dst_dir = os.path.join(src, "Square")


            shutil.move(images[jsq], os.path.join(dst_dir, image_names[jsq]))
        else:
            print(f"Failed to read image: {images[jsq]}")

        progress.update(task, advance=1)
print("处理完毕!")